import pandas as pd
import rqdatac as rq

class Point:
    def __init__(self, dt, point_type, bar):
        self.dt = dt
        self.bar = None
        self.point_type = point_type
        self.point_statue = "temp"
        self.point_value = 0
        self.update_value(bar)
        self.bar_count = 1  # K线数量

    def update_bar(self, bar):
        self.bar_count += 1
        self.update_value(bar)

    def update_value(self, bar):
        self.bar = bar
        self.point_value = bar["high"] if self.point_type == "top" else bar["low"]

    def fix_point(self, bar):
        """将点从Temp状态固定下来"""
        self.update_bar(bar)
        self.point_statue = "new"

    def release_point(self, bar):
        """NEW状态回到Temp状态"""
        self.update_bar(bar)
        self.point_statue = "temp"


class PointManager:
    """
    端点管理器，用来管理使用的点：底分型为的底点，顶分型的为顶点
    点的类型不能变化，即第一次生成是决定该点是什么类型，就只能是这个类型
    点的状态：
        临时：没有出现分型时，最新的点为临时点
        半固定：出分型，但是还有可能会变化，即最新一个分型
        固定：已经固定，不会再有变化的点，一旦固定了，就不能再往回退，即固定的点不允许移动
    """
    def __init__(self, amplitude_rate):
        self.para_point_gap = 5      # 两点之间的K线间隔
        self.para_amplitude_gap = amplitude_rate     # 两点之间的价格幅度间隔
        self.working_point: Point = None
        self.previous_point: Point = None
        self.previous_point_2: Point = None
        self.gap_count = 0  # 上一个点到现在的距离
        self.bar3 = []  # 保存最新的3个bar
        self.left_bar = None    # 保存可以用来给新点使用的3个bar，与 bar3的区别是bar3不区分是否已经使用过
        self.mid_bar = None
        self.right_bar = None

    def update_bar(self, dt, bar):
        # 两种情况会出现3bar没满，一是刚开始数据加载不足，二是上一个分型刚形成，需要重新积累数据
        # 数据不满时，如果新价格超过上一个分型的价格，则要更新上一个分型的位置（低于底分型，高于顶分型）
        # 1、判断是否生成新点
        # 2、判断旧点是否要移动
        # 3、最新点跟随
        self.left_bar = self.mid_bar
        self.mid_bar = self.right_bar
        self.right_bar = bar

        if self.previous_point:
            self.gap_count += 1

        self.bar3.append(bar)
        if len(self.bar3) > 3:
            self.bar3.pop(0)

        # 先判断是否生成新的点

        is_new = self.check_new_point(dt)
        if not is_new:

            # 旧点是否要移动到最新位置（即被同向突破）
            moved = False
            if self.working_point and self.previous_point and self.previous_point_2 and self.previous_point.point_statue != "fixed":
                if self.previous_point.point_type == "top" and bar["high"] > self.previous_point.point_value:
                    self.move_point(bar)
                    moved = True
                elif self.previous_point.point_type == "bot" and bar["low"] < self.previous_point.point_value:
                    self.move_point(bar)
                    moved = True
            # 没有生成新点，不吃点，则移动临时点
            if not moved:
                self.trace_point(dt, bar)
                # 跟踪新点后，判断该点有没有生成新的点，如果没有则按working_point正常更新，如果有则在new_point中发新点
                self.check_new_point(dt)

        if self.previous_point:
            value1 = self.previous_point.point_value
            if self.working_point:
                value2 = self.working_point.point_value
            elif self.previous_point_2:
                value2 = self.previous_point_2.point_value
            else:
                return 0, 0
            close = bar["close"]
            if close < min(value1, value2):
                support = 0
                pressure = (close - min(value1, value2)) * 100 / close
            elif close > max(value1, value2):
                support = (close - max(value1, value2)) * 100 / close
                pressure = 0
            else:
                support = (close - min(value1, value2)) * 100 / close
                pressure = (close - max(value1, value2)) * 100 / close
            return support, pressure
        return 0, 0

    def check_new_point(self, dt):
        if self.left_bar is not None and self.mid_bar is not None and self.right_bar is not None:
            if not self.previous_point and not self.working_point:
                # 生成第一个点
                if self.left_bar["high"] <= self.mid_bar["high"] and self.mid_bar["high"] >= self.right_bar["high"]:
                    # 生成顶点
                    self.previous_point = Point(dt, "top", self.mid_bar)
                    self.previous_point.point_statue = "new"
                    self.working_point = None
                    return True
                if self.left_bar["low"] >= self.mid_bar["low"] and self.mid_bar["low"] <= self.right_bar["low"]:
                    self.previous_point = Point(dt, "bot", self.mid_bar)
                    self.previous_point.point_statue = "new"
                    self.working_point = None
                    return True
            elif self.working_point:
                if self.gap_count < self.para_point_gap:
                    return False
                # 临时点固定
                if self.working_point.point_type == "top":
                    if self.left_bar["high"] <= self.mid_bar["high"] and self.mid_bar["high"] >= self.right_bar["high"]:
                        if abs(self.mid_bar["high"] - self.previous_point.point_value) / self.previous_point.point_value  < self.para_amplitude_gap:
                            return False
                        self.working_point.fix_point(self.mid_bar)
                        #print(self.working_point.dt, self.working_point.point_value)
                        self.previous_point_2 = self.previous_point
                        self.previous_point_2.point_statue = "fixed"
                        self.previous_point = self.working_point
                        self.working_point = None
                        self.gap_count = 0

                        return True
                else:
                    if self.left_bar["low"] >= self.mid_bar["low"] and self.mid_bar["low"] <= self.right_bar["low"]:
                        if abs(self.mid_bar["low"] - self.previous_point.point_value) / self.mid_bar["low"] < self.para_amplitude_gap:
                            return False
                        self.working_point.fix_point(self.mid_bar)
                        #print(self.working_point.dt, self.working_point.point_value)
                        self.previous_point = self.working_point
                        self.previous_point_2 = self.previous_point
                        self.previous_point_2.point_statue = "fixed"
                        self.working_point = None
                        self.gap_count = 0
                        return True
        return False

    def move_point(self, bar):
        """将上一个点重新移到最新的位置
        previous_point要吃掉前面的临时点
        """
        temp_point = self.working_point
        self.working_point = self.previous_point
        self.working_point.bar_count += temp_point.bar_count
        self.gap_count = self.working_point.bar_count
        self.previous_point = self.previous_point_2
        self.previous_point_2 = None
        self.working_point.release_point(bar)
        self.left_bar = self.bar3[0]
        self.mid_bar = self.bar3[1]
        self.right_bar = self.bar3[2]

    def trace_point(self, dt, bar):
        if not self.previous_point:
            return
        if not self.working_point:
            self.working_point = Point(dt, self._get_new_type(), bar)
            self.working_point.point_statue = "temp"
        else:
            self.working_point.update_bar(bar)
    def _get_new_type(self):
        if self.previous_point:
            if self.previous_point.point_type == "top":
                return "bot"
            return "top"
        return "null"


def get_factor_df(sec_id, prices):
    print(sec_id, __file__)
    for k_gap in [3, 6, 15]:
        point_manager = PointManager(float(3/1000))
        supports = []
        pressures = []
        for n, bar in prices.iterrows():
            s, p = point_manager.update_bar(n, bar)
            supports.append(s)
            pressures.append(p)
        prices[f"support_{k_gap}"] = supports
        prices[f"pressure_{k_gap}"] = pressures
    return prices


if __name__ =="__main__":
    if not rq.initialized():
        rq.init("13570866213", "39314656")
    prices = rq.get_price("I2401", start_date="20220801", end_date="20231116", frequency="1m",
                          fields=["open", "close", "high", "low"], expect_df=False)
    point_manager = PointManager(0.003)
    supports = []
    pressures = []
    for n, bar in prices.iterrows():
        s, p = point_manager.update_bar(n, bar)
        supports.append(s)
        pressures.append(p)
